import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime
import os

def create_crime_heatmap(df, save_path=None):
    """
    创建犯罪热力图（使用matplotlib替代folium）
    """
    plt.figure(figsize=(12, 8))
    
    # 创建热力图
    plt.hist2d(df['LON'], df['LAT'], bins=50, cmap='YlOrRd')
    plt.colorbar(label='Crime Count')
    
    plt.title('Crime Heatmap - Los Angeles')
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    
    if save_path:
        # 将.html后缀改为.png
        save_path = save_path.replace('.html', '.png')
        plt.savefig(save_path)
        plt.close()
    
    return plt.gcf()

def plot_crime_by_time(df, save_dir=None):
    """
    绘制犯罪随时间的分布图
    """
    plt.figure(figsize=(15, 10))
    
    # 创建子图
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(20, 15))
    
    # 1. 按小时分布
    hourly_crimes = df['Hour'].value_counts().sort_index()
    ax1.plot(hourly_crimes.index, hourly_crimes.values)
    ax1.set_title('犯罪按小时分布')
    ax1.set_xlabel('小时')
    ax1.set_ylabel('犯罪数量')
    
    # 2. 按星期分布
    daily_crimes = df['DayOfWeek'].value_counts().sort_index()
    ax2.bar(daily_crimes.index, daily_crimes.values)
    ax2.set_title('犯罪按星期分布')
    ax2.set_xlabel('星期')
    ax2.set_ylabel('犯罪数量')
    
    # 3. 按月份分布
    monthly_crimes = df['Month'].value_counts().sort_index()
    ax3.plot(monthly_crimes.index, monthly_crimes.values, marker='o')
    ax3.set_title('犯罪按月份分布')
    ax3.set_xlabel('月份')
    ax3.set_ylabel('犯罪数量')
    
    # 4. 按季节分布
    seasonal_crimes = df['Season'].value_counts()
    ax4.pie(seasonal_crimes.values, labels=seasonal_crimes.index, autopct='%1.1f%%')
    ax4.set_title('犯罪按季节分布')
    
    plt.tight_layout()
    
    if save_dir:
        plt.savefig(os.path.join(save_dir, 'crime_time_distribution.png'))
        plt.close()
    
    return fig

def plot_crime_types(df, top_n=10, save_dir=None):
    """
    绘制犯罪类型分布图
    """
    plt.figure(figsize=(15, 8))
    
    # 获取前N种最常见的犯罪类型
    crime_counts = df['Crm Cd Desc'].value_counts().head(top_n)
    
    # 创建水平条形图
    plt.barh(crime_counts.index, crime_counts.values)
    plt.title(f'前{top_n}种最常见犯罪类型')
    plt.xlabel('犯罪数量')
    
    plt.tight_layout()
    
    if save_dir:
        plt.savefig(os.path.join(save_dir, 'crime_types_distribution.png'))
        plt.close()
    
    return plt.gcf()

def plot_area_crime_density(df, save_dir=None):
    """
    绘制区域犯罪密度图
    """
    plt.figure(figsize=(12, 8))
    
    # 计算每个区域的犯罪密度
    area_density = df.groupby('AREA NAME').size().sort_values(ascending=True)
    
    # 创建水平条形图
    plt.barh(area_density.index, area_density.values)
    plt.title('各区域犯罪密度')
    plt.xlabel('犯罪数量')
    
    plt.tight_layout()
    
    if save_dir:
        plt.savefig(os.path.join(save_dir, 'area_crime_density.png'))
        plt.close()
    
    return plt.gcf()

def visualize_hotspots(df, hotspots, probabilities, save_path=None):
    """
    可视化预测的犯罪热点（使用matplotlib替代folium）
    """
    plt.figure(figsize=(12, 8))
    
    # 绘制所有点
    plt.scatter(df['LON'][~hotspots], df['LAT'][~hotspots], 
               c='blue', alpha=0.5, s=10, label='Normal')
    plt.scatter(df['LON'][hotspots], df['LAT'][hotspots],
               c='red', alpha=0.5, s=10, label='Hotspot')
    
    plt.title('Predicted Crime Hotspots - Los Angeles')
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    plt.legend()
    
    if save_path:
        # 将.html后缀改为.png
        save_path = save_path.replace('.html', '.png')
        plt.savefig(save_path)
        plt.close()
    
    return plt.gcf()


def create_visualization_report(df, model_results, save_dir="data/visualizations"):
    """
    创建完整的可视化报告
    """
    # 获取绝对路径并创建保存目录
    save_dir = os.path.abspath(save_dir)
    os.makedirs(save_dir, exist_ok=True)

    print(f"所有图像将保存到：{save_dir}")

    # 1. 创建热力图
    print("Creating heatmap...")
    create_crime_heatmap(df, os.path.join(save_dir, 'crime_heatmap.png'))

    # 2. 创建时间分布图
    print("Creating time distribution plots...")
    plot_crime_by_time(df, save_dir)

    # 3. 创建犯罪类型分布图
    print("Creating crime type distribution...")
    plot_crime_types(df, save_dir=save_dir)

    # 4. 创建区域密度图
    print("Creating area density plot...")
    plot_area_crime_density(df, save_dir=save_dir)

    # 5. 创建模型评估可视化
    if model_results:
        print("Creating model evaluation plots...")
        plt.figure(figsize=(10, 6))
        model_scores = {name: results['test_score']
                        for name, results in model_results.items()}
        plt.bar(model_scores.keys(), model_scores.values())
        plt.title('模型性能比较')
        plt.ylabel('测试集准确率')
        plt.xticks(rotation=45)
        plt.tight_layout()
        plt.savefig(os.path.join(save_dir, 'model_comparison.png'))
        plt.close()

    print(f"✅ 可视化报告已生成：{save_dir}")


if __name__ == "__main__":
    # 加载数据
    print("Loading data...")
    df = pd.read_csv("data/processed/processed_crime_data.csv")
    
    # 创建可视化
    print("Creating visualizations...")
    create_visualization_report(df, None)
    
    print("Visualization completed successfully!") 